/* 
 * Copyright 2014 Josef Hertl.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package caetus.pkcs11.key;

import caetus.pkcs11.template.Template;
import caetus.pkcs11.template.TemplateMaker;
import java.util.Random;
import sun.security.pkcs11.wrapper.CK_ATTRIBUTE;
import sun.security.pkcs11.wrapper.CK_MECHANISM;
import sun.security.pkcs11.wrapper.PKCS11;
import sun.security.pkcs11.wrapper.PKCS11Constants;
import sun.security.pkcs11.wrapper.PKCS11Exception;

/**
 * Represents an abstract key. Splits further into symmetric and asymmetric.
 *
 * @author Josef Hertl
 */
public abstract class Key {

    public Key(long sessionHandle, PKCS11 pkcs11, CipherEnum cipherEnum) {
        this.sessionHandle = sessionHandle;
        this.pkcs11 = pkcs11;
        this.cipherEnum = cipherEnum;
    }
    /**
     * The pointer to the location in the token that holds the key.
     */
    protected long keyHandle;
    /**
     * The handle of the session this key runs in.
     */
    protected long sessionHandle;
    protected PKCS11 pkcs11;
    protected CipherEnum cipherEnum;
    /**
     * Determines if the key has a value assigned, so that it can perform
     * cryptographic operations.
     */
    private boolean initialized = false;
    /**
     * Just for convenience.
     */
    protected TemplateMaker templateMaker = new TemplateMaker();
    /**
     * A random for convenience.
     */
    protected Random random = new Random();

    /**
     * Uses the key to encrypt the data.
     *
     * @param dataToEncrypt The data that will be encrypted. The length does not
     * matter.
     * @return Encrypted data.
     */
    public abstract byte[] encrypt(byte[] dataToEncrypt) throws PKCS11Exception;

    /**
     * Uses the key to decrypt input data.
     *
     * @param dataToDecrypt The data to be decrypted.
     * @return The decrypted version of the input data.
     */
    public abstract byte[] decrypt(byte[] dataToDecrypt) throws PKCS11Exception;

    /**
     * Wraps another key.
     *
     * @param keyToWrap The key that will be wrapped using this key.
     * @return The newly wrapped key in the form of byte[]
     */
    public abstract byte[] wrapKey(Key keyToWrap) throws PKCS11Exception;

    /**
     * Unwraps another secret key.
     *
     * @param toUnwrap The bytes from which the key should be unwrapped.
     * @param template The template that is assigned to the unwrapped key.
     * @param cipher The cipher of this key.
     * @return The newly created Key object.
     */
    public abstract SecretKey unwrapSecretKey(byte[] toUnwrap, Template template, CipherEnum cipher) throws PKCS11Exception;

    /**
     * Unwraps another private key.
     *
     * @param toUnwrap The bytes from which the key should be unwrapped.
     * @param template The template that is assigned to the unwrapped key.
     * @return The newly created Key object.
     */
    public abstract PrivateKey unwrapPrivateKey(byte[] toUnwrap, Template template, PublicKey correspondingPublicKey) throws PKCS11Exception;

    /**
     * @return True if the key can perform cryptographic operations, false
     * otherwise.
     */
    public boolean isInitialized() {
        return initialized;
    }

    /**
     * Marks the key as initialized, so that it is capable of performing
     * cryptographic operations.
     */
    protected void initialize() {
        initialized = true;
    }

    /**
     * @return Handle of the key - pointer to where it is saved on the token.
     */
    public long getKeyHandle() {
        return keyHandle;
    }

    /**
     * Sets the key handle and also initializes the key.
     */
    public void setKeyHandleAndInitialize(long keyHandle) {
        this.keyHandle = keyHandle;
        initialize();
    }

    public CipherEnum getCipherEnum() {
        return cipherEnum;
    }

    public void setCipherEnum(CipherEnum cipherEnum) {
        this.cipherEnum = cipherEnum;
    }

    /**
     * Sets the given attribute.
     *
     * @param ck_attribute The attribute to set.
     */
    public void setAttribute(CK_ATTRIBUTE ck_attribute) throws PKCS11Exception {
        CK_ATTRIBUTE[] setAttributes = new CK_ATTRIBUTE[1];
        setAttributes[0] = ck_attribute;
        pkcs11.C_SetAttributeValue(sessionHandle, keyHandle, setAttributes);
    }

    /**
     * Changes multiple attributes, as specified by the given template.
     *
     * @param template The template for attribute change.
     */
    public void setAttributes(Template template) throws PKCS11Exception {
        pkcs11.C_SetAttributeValue(sessionHandle, keyHandle, template.finishAndGetAttributes());
    }

    /**
     * Gets the value of the given attribute.
     *
     * @param attribute The desired attribute.
     * @return The single attribute in plain form.
     */
    public CK_ATTRIBUTE getAttribute(CK_ATTRIBUTE attribute) throws PKCS11Exception {
        try {
            CK_ATTRIBUTE[] resultAttributes = new CK_ATTRIBUTE[1];
            resultAttributes[0] = attribute;
            pkcs11.C_GetAttributeValue(sessionHandle, keyHandle, resultAttributes);
            return resultAttributes[0];
        } catch (OutOfMemoryError error) {
            throw new PKCS11Exception(PKCS11Constants.CKR_GENERAL_ERROR);
        }
    }

    /**
     * @return True if the given attribute is true, false otherwise.
     */
    public boolean attributeIsTrue(CK_ATTRIBUTE ck_attribute) throws PKCS11Exception {
        CK_ATTRIBUTE[] resultAttributes = new CK_ATTRIBUTE[1];
        resultAttributes[0] = ck_attribute;
        pkcs11.C_GetAttributeValue(sessionHandle, keyHandle, resultAttributes);
        return resultAttributes[0].getBoolean();
    }

    /**
     * @return True if the key is sensitive, false otherwise.
     */
    public boolean isSensitive() throws PKCS11Exception {
        CK_ATTRIBUTE[] resultAttributes = new CK_ATTRIBUTE[1];
        resultAttributes[0] = new CK_ATTRIBUTE(PKCS11Constants.CKA_SENSITIVE);
        pkcs11.C_GetAttributeValue(sessionHandle, keyHandle, resultAttributes);
        return resultAttributes[0].getBoolean();
    }

    /**
     * The key copies itself.
     *
     * @param template The template to pass to the copied key.
     * @return copy of this key with the given template.
     */
    public abstract Key copyObject(Template template) throws PKCS11Exception;

    /**
     * The key derives a new key based on itself and the input mechanism. Other
     * input is randomly generated based on the mechanism.
     *
     * @param mechanism The mechanism to use - CKM_CONCATENATE_BASE_AND_KEY,
     * CKM_CONCATENATE_BASE_AND_DATA, CKM_CONCATENATE_DATA_AND_BASE,
     * CKM_XOR_BASE_AND_DATA, CKM_EXTRACT_KEY_FROM_KEY
     * @param template The template for the newly derived key.
     * @return The newly derived key.
     */
    public abstract Key deriveKey(CK_MECHANISM mechanism, Template template) throws PKCS11Exception;
}
